# -*- coding: utf-8 -*-
# @Time    : 2024/9/13 20:49
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : particle_initializer.py
# @Software: PyCharm
import enum

import matplotlib

import mot.atom

matplotlib.use('tkagg')
import matplotlib.pyplot as plt
import numpy
import pandas
from scipy.stats import maxwell
from myutils.test_utils import timer
import scipy.constants as C
from _logging import logger
logger.info("I am reloaded")

class ParticleInitializerBase:
    class Columns(enum.Enum):
        x = 0
        y = 1
        z = 2
        vx = 10
        vy = 11
        vz = 12

    columns = [col.name for col in Columns]

    def generate(self) -> pandas.DataFrame:
        pass


class MaxwellParticleInitializer(ParticleInitializerBase):
    def __init__(self, n0, xyz_range: numpy.ndarray, vmin, vmax, T0, m):
        """
        :param n0:
        :param vmin:
        :param vmax:
        :param xyz_range: shape (3, 2).
        :param m: mass of particle, unit in kg.
        """
        self.n0 = n0
        self.vmin = vmin
        self.vmax = vmax
        self.xyz_range = xyz_range
        self.T0 = T0
        self.m = m

    @timer
    def generate(self) -> pandas.DataFrame:
        pos: numpy.ndarray = self.xyz_range[:, 0] + numpy.random.random((self.n0, 3)) * (
                    self.xyz_range[:, 1] - self.xyz_range[:, 0])
        thetas_phis = numpy.random.random((self.n0, 2))*2*numpy.pi
        v_abs = maxwell.rvs(size=pos.shape[0]) / (self.m / (C.k * self.T0))**0.5
        index_=numpy.where((v_abs<self.vmax)&(v_abs>self.vmin))
        vel =( v_abs* (numpy.array(
            [ numpy.sin(thetas_phis[:, 1]) * numpy.cos(thetas_phis[:, 0]),
              numpy.sin(thetas_phis[:, 1]) *  numpy.sin(thetas_phis[:, 0]),
              numpy.cos(thetas_phis[:, 1]), ]))).T

        logger.info("max velocity = %.2e m/s (截断前)"%(v_abs.max()))

        return pandas.DataFrame(numpy.array([*pos.T, *vel.T]).T[index_[0]], columns=self.columns)


if __name__ == '__main__':
    plt.ion()
    par_ini = MaxwellParticleInitializer(int(1e4),
                                         numpy.array([[-2e-3, 2e-3, ],
                                                      [-2e-3, 3e-3, ],
                                                      [-2e-3, 4e-3, ],
                                                      ]), 0, 20, 1.0, mot.atom.Rb87_D2.mass_kg)
    par_data = par_ini.generate()
    plt.figure()
    plt.scatter(*par_data[['x', 'y']].values.T)

    plt.figure()
    vs= numpy.linalg.norm(par_data[['vx', 'vy', 'vz']].values, axis=1)
    plt.hist(numpy.linalg.norm(par_data[['vx','vy','vz']].values,axis =  1), 100, density=True)
    _vs = numpy.linspace(vs.min()*0, vs.max(), 100)


    plt.figure()
    plt.scatter(par_data['x'],par_data['vx'])

